#!/bin/bash

# SSH密钥配置脚本
# 用于配置服务器SSH密钥登录

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

print_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 检查本地SSH密钥
check_local_keys() {
    print_info "检查本地SSH密钥..."
    
    if [ -f ~/.ssh/id_rsa.pub ]; then
        print_success "找到SSH公钥: ~/.ssh/id_rsa.pub"
        return 0
    else
        print_error "未找到SSH公钥，请先生成SSH密钥对："
        echo "ssh-keygen -t rsa -b 4096 -C 'your-email@example.com'"
        return 1
    fi
}

# 获取服务器密码
get_server_password() {
    echo ""
    print_warning "需要服务器root密码来配置SSH密钥"
    read -s -p "请输入服务器root密码: " SERVER_PASSWORD
    echo ""
}

# 配置服务器SSH密钥
setup_server_ssh() {
    print_info "配置服务器SSH密钥..."
    
    # 读取本地公钥
    LOCAL_PUBLIC_KEY=$(cat ~/.ssh/id_rsa.pub)
    
    # 使用密码登录服务器并配置SSH密钥
    sshpass -p "$SERVER_PASSWORD" ssh -o StrictHostKeyChecking=no root@47.120.11.77 << 'EOF'
        # 创建.ssh目录
        mkdir -p ~/.ssh
        chmod 700 ~/.ssh
        
        # 创建authorized_keys文件
        touch ~/.ssh/authorized_keys
        chmod 600 ~/.ssh/authorized_keys
        
        # 备份现有配置
        if [ -f ~/.ssh/authorized_keys ]; then
            cp ~/.ssh/authorized_keys ~/.ssh/authorized_keys.backup
        fi
EOF
    
    # 上传公钥到服务器
    echo "$LOCAL_PUBLIC_KEY" | sshpass -p "$SERVER_PASSWORD" ssh -o StrictHostKeyChecking=no root@47.120.11.77 "cat >> ~/.ssh/authorized_keys"
    
    print_success "SSH密钥配置完成"
}

# 测试SSH密钥登录
test_ssh_key() {
    print_info "测试SSH密钥登录..."
    
    # 等待一下让服务器处理
    sleep 2
    
    if ssh -o ConnectTimeout=10 -o BatchMode=yes feishu-server "echo 'SSH密钥登录测试成功'" 2>/dev/null; then
        print_success "SSH密钥登录测试成功"
        return 0
    else
        print_error "SSH密钥登录测试失败"
        return 1
    fi
}

# 安装sshpass（如果需要）
install_sshpass() {
    if ! command -v sshpass &> /dev/null; then
        print_info "安装sshpass..."
        
        if [[ "$OSTYPE" == "darwin"* ]]; then
            # macOS
            if command -v brew &> /dev/null; then
                brew install sshpass
            else
                print_error "请先安装Homebrew，然后运行: brew install sshpass"
                return 1
            fi
        elif [[ "$OSTYPE" == "linux-gnu"* ]]; then
            # Linux
            if command -v apt-get &> /dev/null; then
                sudo apt-get update && sudo apt-get install -y sshpass
            elif command -v yum &> /dev/null; then
                sudo yum install -y sshpass
            else
                print_error "无法自动安装sshpass，请手动安装"
                return 1
            fi
        else
            print_error "不支持的操作系统: $OSTYPE"
            return 1
        fi
    fi
}

# 显示配置信息
show_config_info() {
    echo ""
    echo "🔑 SSH密钥配置信息："
    echo "   服务器IP: 47.120.11.77"
    echo "   用户名: root"
    echo "   SSH别名: feishu-server"
    echo "   公钥文件: ~/.ssh/id_rsa.pub"
    echo ""
    echo "📋 测试命令："
    echo "   ssh feishu-server"
    echo "   ssh feishu-server 'echo 连接成功'"
    echo ""
    echo "🔧 如果配置失败，请手动执行："
    echo "   1. ssh root@47.120.11.77"
    echo "   2. mkdir -p ~/.ssh"
    echo "   3. chmod 700 ~/.ssh"
    echo "   4. touch ~/.ssh/authorized_keys"
    echo "   5. chmod 600 ~/.ssh/authorized_keys"
    echo "   6. 将以下公钥内容添加到 ~/.ssh/authorized_keys："
    echo ""
    cat ~/.ssh/id_rsa.pub
    echo ""
}

# 主函数
main() {
    echo "🔑 SSH密钥配置脚本"
    echo "服务器: 47.120.11.77"
    echo ""
    
    # 检查本地SSH密钥
    if ! check_local_keys; then
        exit 1
    fi
    
    # 安装sshpass
    if ! install_sshpass; then
        exit 1
    fi
    
    # 获取服务器密码
    get_server_password
    
    # 配置服务器SSH密钥
    setup_server_ssh
    
    # 测试SSH密钥登录
    if test_ssh_key; then
        print_success "SSH密钥配置成功！"
        show_config_info
    else
        print_error "SSH密钥配置失败，请检查："
        echo "1. 服务器密码是否正确"
        echo "2. 服务器是否允许密码登录"
        echo "3. 防火墙是否开放SSH端口"
        echo ""
        show_config_info
    fi
}

# 运行主函数
main "$@" 